const tf = require('@tensorflow/tfjs-node');
const mobilenet = require('@tensorflow-models/mobilenet'); // MobileNet 模型，用于迁移学习
const fs = require('fs-extra'); // 一个扩展了 Node.js 文件系统模块的库，用于文件操作

// 定义模型结构
const numClasses = 3;
const model = tf.sequential();
// model.add(tf.layers.flatten({ inputShape: [7, 7, 256] }));
model.add(tf.layers.flatten({ inputShape: [96, 96, 3] }));
model.add(tf.layers.dense({ units: 100, activation: 'relu' }));
model.add(tf.layers.dense({ units: numClasses, activation: 'softmax' }));

const labelArray = {
    daisy: 0,
    dandelion: 1,
    rose: 2
}

// 编译模型
model.compile({
  optimizer: tf.train.adam(0.0001),
  loss: 'categoricalCrossentropy',
  metrics: ['accuracy'],
});
const trainDataPath = '../image/flowers';
const testDataPath = '../image/test';
let imagesTensor;
// 从文件夹中读取图片数据
async function getFileListAndLabel(path) {
    const imageData = []
    const files = await fs.readdirSync(path)
    const images = files.reduce((prev, dir) => {
        let fileList = fs.readdirSync(path + '/' + dir)
        fileList.map(file => {
            // labels.push(dir)
            // return `${path}/${dir}/${file}`
            imageData.push({
                imagePath: `${path}/${dir}/${file}`,
                label: labelArray[dir],
                className: dir
            })
        })
    }, [])
    return imageData
}
// 图片路径数据读取图片数据并转化为张量
const imageToTensor = (paths) => {
  // concat方法是将多个张量合并成一个张量
  return tf.concat(paths.map(path => {
    // 使用tidy方法，对张量进行内存控制，不然你直接读直接内存炸
    try {
        return tf.tidy(() => {
          return tf.node
          .decodeImage(fs.readFileSync(path.imagePath)) // 把图片转换成张量
          .resizeNearestNeighbor([96, 96]) // 把图片的宽高转换成96x96
          .toFloat() // 转换成浮点型的张量
          .div(tf.scalar(255.0)) // 将图片张量进行归一化（这个自行去理解）
          .expandDims()
        })
    } catch {
        console.log('path', path.imagePath)
    }
  }))
}
// 从文件中加载训练数据
async function loadTrainData(dirPath) {
  try {
    const data = await getFileListAndLabel(dirPath)
    // tf.util.shuffle(data)
    console.log(data)
    // const rawData = await fs.readFile(trainDataPath, 'utf8');
    // const data = JSON.parse(rawData);

    // 处理数据，将图像数据转换为 TensorFlow.js 的 Tensor
    // const images = data.map(item => {
    //   const imageBuffer = fs.readFileSync(item.imagePath);
    //   const tfimage = tf.node.decodeImage(imageBuffer);
    //   return tfimage;
    // });
    imagesTensor = imageToTensor(data)

    // 将标签转换为 TensorFlow.js 的 Tensor
    const labels = tf.tensor1d(data.map(item => item.label), 'int32');
    const classNames = data.map(item => item.className);
    // 创建 tf.data.Dataset
    // const dataset = tf.data.zip({ images: images, labels: labels });
    const dataset = { images: imagesTensor, labels };
    console.log('dataset', dataset)
    // 返回 Dataset
    return dataset
  } catch (error) {
    console.error('加载训练数据时出错:', error);
    throw error;
  }
}
// loadTrainData(trainDataPath)

async function start() {
    // 准备数据
    const trainData = await loadTrainData(trainDataPath)/* 从文件中加载训练数据 */;
    const testData = await loadTrainData(testDataPath)/* 从文件中加载测试数据 */;
    // console.log(trainData.labels)
    // 将标签转换为独热编码
    const trainLabels = tf.oneHot(trainData.labels, numClasses);
    const testLabels = tf.oneHot(testData.labels, numClasses);
    
    // 训练模型
    model.fit(trainData.images, trainLabels, {
      epochs: 15,
      validationData: [testData.images, testLabels],
      shuffle: true, // 打乱数据
    }).then(info => {
      console.log('训练完成:', info);
      // 保存模型
      model.save('file://model').then(() => {
        console.log('模型保存成功。');
      });
    });
}
start()